import tensorflow as tf
import time
import grpc
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc

def main(_):
    # 设置grpc
    options = [('grpc.max_send_message_length', 1000 * 1024 * 1024),
            ('grpc.max_receive_message_length', 1000 * 1024 * 1024)]
    channel = grpc.insecure_channel('139.199.34.94:8500', options = options)
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'half_plus'
    request.model_spec.signature_name = tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY

    tensor_x = tf.compat.v1.make_tensor_proto([2.0,4.0],dtype=tf.float32,shape=[2])
    tensor_y = tf.compat.v1.make_tensor_proto([3.0, 6.0], dtype=tf.float32, shape=[2])
    init = tf.global_variables_initializer()

    request.inputs['input_x'].CopyFrom(tensor_x)
    request.inputs['input_y'].CopyFrom(tensor_y)
    start = time.time()

    # 法一，速度较慢
    # result = stub.Predict(request, 10.0)  # 10 secs timeout

    # 法二，速度较快
    result_future = stub.Predict.future(request, 30.0)  # 10 secs timeout

    response = result_future.result()

    results = {}
    for key in response.outputs:
        tensor_proto = response.outputs[key]
        nd_array = tf.contrib.util.make_ndarray(tensor_proto)
        results[key] = nd_array
    print("cost %ss to predict: " % (time.time() - start))
    print("predict label is:", results)
    stop = time.time()
    print('time is ', stop - start)
    # print(type(result))
    # print( result.outputs)


if __name__ == '__main__':
  tf.app.run()